package com.mask.encryption.service.impl;

import com.mask.encryption.service.MaskEncryptionService;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.security.KeyFactory;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;

/**
 * 混合加密：RSA(OAEP-SHA256) 加密随机 AES-GCM 密钥，数据用 AES-GCM 加密
 * 密文结构： [RSA_KEY_LEN(2B)] [RSA(AES_KEY)] [IV(12B)] [AES_CIPHER]
 */
public class HybridRsaAesEncryptionService implements MaskEncryptionService {

    private static final String RSA_TRANSFORMATION = "RSA/ECB/OAEPWithSHA-256AndMGF1Padding";
    private static final String AES_TRANSFORMATION = "AES/GCM/NoPadding";
    private static final int IV_LENGTH = 12;

    private final PublicKey publicKey;
    private final PrivateKey privateKey;
    private final int gcmTagLength;
    private final SecureRandom secureRandom = new SecureRandom();

    public HybridRsaAesEncryptionService(String publicKeyPem, String privateKeyPem, int gcmTagLength) {
        try {
            this.publicKey = parsePublicKey(publicKeyPem);
            this.privateKey = parsePrivateKey(privateKeyPem);
            this.gcmTagLength = gcmTagLength;
        } catch (Exception e) {
            throw new IllegalArgumentException("Invalid RSA keys", e);
        }
    }

    private static PublicKey parsePublicKey(String pem) throws Exception {
        String normalized = pem.replace("-----BEGIN PUBLIC KEY-----", "")
                .replace("-----END PUBLIC KEY-----", "")
                .replaceAll("\\s", "");
        byte[] keyBytes = Base64.getDecoder().decode(normalized);
        X509EncodedKeySpec spec = new X509EncodedKeySpec(keyBytes);
        return KeyFactory.getInstance("RSA").generatePublic(spec);
    }

    private static PrivateKey parsePrivateKey(String pem) throws Exception {
        String normalized = pem.replace("-----BEGIN PRIVATE KEY-----", "")
                .replace("-----END PRIVATE KEY-----", "")
                .replaceAll("\\s", "");
        byte[] keyBytes = Base64.getDecoder().decode(normalized);
        PKCS8EncodedKeySpec spec = new PKCS8EncodedKeySpec(keyBytes);
        return KeyFactory.getInstance("RSA").generatePrivate(spec);
    }

    @Override
    public byte[] encrypt(byte[] plaintext) {
        try {
            // 1) 生成临时 AES 密钥
            KeyGenerator kg = KeyGenerator.getInstance("AES");
            kg.init(256);
            SecretKey aesKey = kg.generateKey();

            // 2) AES-GCM 加密数据，随机 IV
            byte[] iv = new byte[IV_LENGTH];
            secureRandom.nextBytes(iv);
            Cipher aes = Cipher.getInstance(AES_TRANSFORMATION);
            aes.init(Cipher.ENCRYPT_MODE, aesKey, new GCMParameterSpec(gcmTagLength, iv));
            byte[] aesCipher = aes.doFinal(plaintext);

            // 3) 使用 RSA 公钥加密 AES 密钥
            Cipher rsa = Cipher.getInstance(RSA_TRANSFORMATION);
            rsa.init(Cipher.ENCRYPT_MODE, publicKey);
            byte[] rsaEncryptedKey = rsa.doFinal(aesKey.getEncoded());

            // 4) 组装密文：2B 长度（大端） + RSA密钥密文 + IV + AES密文
            if (rsaEncryptedKey.length > 65535) {
                throw new IllegalStateException("RSA encrypted key too long");
            }
            ByteBuffer buffer = ByteBuffer.allocate(2 + rsaEncryptedKey.length + IV_LENGTH + aesCipher.length);
            buffer.putShort((short) (rsaEncryptedKey.length & 0xFFFF));
            buffer.put(rsaEncryptedKey);
            buffer.put(iv);
            buffer.put(aesCipher);
            return buffer.array();
        } catch (Exception e) {
            throw new IllegalStateException("Hybrid encrypt failed", e);
        }
    }

    @Override
    public byte[] decrypt(byte[] ciphertext) {
        try {
            ByteBuffer buffer = ByteBuffer.wrap(ciphertext);
            int rsaLen = Short.toUnsignedInt(buffer.getShort());
            byte[] rsaEncryptedKey = new byte[rsaLen];
            buffer.get(rsaEncryptedKey);
            byte[] iv = new byte[IV_LENGTH];
            buffer.get(iv);
            byte[] aesCipher = new byte[buffer.remaining()];
            buffer.get(aesCipher);

            // 1) 用 RSA 私钥解密 AES key
            Cipher rsa = Cipher.getInstance(RSA_TRANSFORMATION);
            rsa.init(Cipher.DECRYPT_MODE, privateKey);
            byte[] aesKeyBytes = rsa.doFinal(rsaEncryptedKey);
            SecretKeySpec aesKey = new SecretKeySpec(aesKeyBytes, "AES");

            // 2) 用 AES-GCM 解密数据
            Cipher aes = Cipher.getInstance(AES_TRANSFORMATION);
            aes.init(Cipher.DECRYPT_MODE, aesKey, new GCMParameterSpec(gcmTagLength, iv));
            return aes.doFinal(aesCipher);
        } catch (Exception e) {
            throw new IllegalStateException("Hybrid decrypt failed", e);
        }
    }

    @Override
    public String encryptToBase64(String plaintext) {
        return Base64.getEncoder().encodeToString(encrypt(plaintext.getBytes(StandardCharsets.UTF_8)));
    }

    @Override
    public String decryptFromBase64(String base64Ciphertext) {
        byte[] cipher = Base64.getDecoder().decode(base64Ciphertext);
        return new String(decrypt(cipher), StandardCharsets.UTF_8);
    }
}


